#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
将 json + jsonl 转换为统一检测接口格式
=====================================

【功能】
- 输入一个 JSON（含 image_path 信息）和一个 JSONL（含模型 predict 标签）
- 顺序对应检查：条目数必须一致
- 从 JSON 中取出 image_path
- 从 JSONL 的 predict 字段中解析 <region> 与 <description>
- 输出目标 JSON，格式为：
  [
    {
      "image_path": "/path/to/image1.png",
      "bboxes": [
        {"xyxy": [100, 200, 400, 500], "description": "出血点？可能位于黄斑旁"},
        {"xyxy": [50, 80, 120, 140], "description": "微动脉瘤"}
      ]
    },
    ...
  ]

【用法示例】
python ./api/sft2checkZoneAPI.py  \
  --dataset_file /home/zhangpinglu/data0/gy/code/fundus-reasoner-adaptive/experiments/dataset/Alpaca_data/mix_cot_only_test.json \
  --prediction_file  /home/zhangpinglu/data0/gy/code/fundus-reasoner-adaptive/experiments/eval/colda_start/sft_stage2_mix_checkpoint-400.jsonl --output_file ./experiments/mix_cot_only.json

作者：zym1105
日期：2025-08-28
"""

import argparse
import json,os
import re

def parse_args():
    parser = argparse.ArgumentParser(description="Convert json + jsonl to bbox format json")
    parser.add_argument(
        "--dataset_file",
        type=str,
        default="experiments/dataset/Alpaca_data/mix_cot_only_test.json",
        help="输入的 JSON 文件（包含 image_path）"
    )
    parser.add_argument(
        "--prediction_file",
        type=str,
        default="experiments/eval/colda_start/sft_stage2_mix_checkpoint-400.jsonl",
        help="输入的 JSONL 文件（包含 predict 输出）"
    )
    parser.add_argument(
        "--output_file",
        type=str,
        default="experiments/eval/final_eval_input.json",
        help="输出 JSON 文件"
    )
    return parser.parse_args()

def extract_bboxes_from_predict(predict: str):
    """
    从 predict 文本中提取所有 region + description 对
    """
    bboxes = []
    if not predict:
        return bboxes

    # 正则提取所有 region/description 对
    pattern = re.compile(
        r"<region>\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)</region>\s*<description>(.*?)</description>",
        re.S
    )
    matches = pattern.findall(predict)
    for m in matches:
        x1, y1, x2, y2, desc = m
        bboxes.append({
            "xyxy": [int(x1), int(y1), int(x2), int(y2)],
            "description": desc.strip()
        })
    return bboxes

def main():
    args = parse_args()

    # 读取 json 和 jsonl
    with open(args.dataset_file, "r", encoding="utf-8") as f:
        json_data = json.load(f)
    with open(args.prediction_file, "r", encoding="utf-8") as f:
        jsonl_data = [json.loads(line) for line in f]

    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    # 检查条数是否一致
    if len(json_data) != len(jsonl_data):
        raise ValueError(f"条目数不一致: json={len(json_data)}, jsonl={len(jsonl_data)}")

    results = []
    for item_json, item_jsonl in zip(json_data, jsonl_data):
        image_path = item_json["images"][0] if "images" in item_json and item_json["images"] else None
        predict = item_jsonl.get("predict", "")

        bboxes = extract_bboxes_from_predict(predict)
        results.append({
            "image_path": image_path,
            "bboxes": bboxes
        })

    with open(args.output_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

    print(f"[SUCCESS] 已保存到 {args.output_file}")

if __name__ == "__main__":
    main()
